package egoc

import (
	"fmt"
	"net/http"

	"golang.org/x/net/websocket"
)

// CheckWebsocket 升级为 websocket 连接
func (c *Context) CheckWebsocket(handle websocket.Handler, check func(*websocket.Config, *http.Request) error) {
	websocket.Server{
		Handler:   handle,
		Handshake: check,
	}.ServeHTTP(c.Writer, c.Request)
}

// UpgradeWebsocket 升级为 websocket 连接
// checkOrigin 是否检查 origin 头
func (c *Context) UpgradeWebsocket(handle websocket.Handler, checkOrigin bool) {
	websocket.Server{Handler: handle,
		Handshake: func(config *websocket.Config, req *http.Request) error {
			var err error
			config.Origin, err = websocket.Origin(config, req)
			if checkOrigin {
				if err == nil && config.Origin == nil {
					return fmt.Errorf("null origin")
				}
			}
			return err
		},
	}.ServeHTTP(c.Writer, c.Request)
}
